
# third party
import socket
import struct
import Queue
import time
from threading import Thread
from threading import Lock
from threading import Timer

# first party
from misc import threadPrint


# peer protocol message IDs
CHOKE = 0
UNCHOKE = 1
INTERESTED = 2
NOT_INTERESTED = 3
HAVE = 4
BITFIELD = 5
REQUEST = 6
FRAGMENT = 7
CANCEL = 8
PORT = 9

ACTIVE_REQUESTS_ALLOWED = 10


class Peer (Thread):
	def __init__ (self, torrent, ip, port):
		Thread.__init__(self)
		
		# basics
		self.alive = True
		self.torrent = torrent
		self.ip = ip
		self.port = port
		
		# the send/recv machinery uses these		
		self.sendThrottle = torrent.owner.sendThrottle
		self.recvThrottle = torrent.owner.recvThrottle

		self.recvData = ""
		self.recvSize = 0
		self.recvTotal = 0

		self.sendQueue = Queue.Queue()
		self.sendData = ""
		self.sendTotal = 0

		self.activeRequests = []
		self.queuedRequests = []
		
		self.fragmentsReceived = 0
		self.fragmentsSent = 0

		# timers that keep me alive and kill me, respectively
		self.keepAliveTimer = None
		self.timeout = None
	
		# my stance towards the peer
		self.choked = True
		self.interesting = False

		# the peer's stance towards me
		self.choking_me = True
		self.interested_in_me = False

		# a list of pieces available
		self.pieceComplete = []
		for i in range(self.torrent.pieceCount):
			self.pieceComplete.append( False )

		self.start()

	def statusString (self):
		string = ""
		string += self.ip.rjust(15)
		string += str(self.sendTotal / 1024).rjust(9)
		string += str(self.recvTotal / 1024).rjust(10)
		string += str(self.fragmentsSent).rjust(8)
		string += str(self.fragmentsReceived).rjust(10)

		if self.choking_me:
			string += " Choking"
		else:
			string += "        "

		if self.interested_in_me:
			string += " Interested"
		else:
			string += "           "

		return string


	def initFromTracker (self):
		self.sock = socket.socket()
		self.sock.setblocking(1)
		
		try:
			self.sock.connect( (self.ip, self.port) )
			self.sock.send( self.torrent.handshake )

			self.sendBitfield()
			self.setChoke(False)

			pstrlen = ""
			while len(pstrlen) == 0:
				pstrlen = self.sock.recv(1)
			pstrlen = struct.unpack( "!B", pstrlen )[0]

			pstr = ""
			while len(pstr) < pstrlen:
				pstr += self.sock.recv( pstrlen - len(pstr) )
		
			if pstr != "BitTorrent protocol":
				raise Exception("ERROR: unexpected protocol string :" + pstr)

			reservedBytes = ""
			while len(reservedBytes) < 8:
				reservedBytes += self.sock.recv( 8 - len(reservedBytes) )

			torrentGUID = ""
			while len(torrentGUID) < 20:
				torrentGUID += self.sock.recv( 20 - len(torrentGUID) )

			self.id = ""
			while len(self.id) < 20:
				self.id += self.sock.recv( 20 - len(self.id) )

			if torrentGUID != self.torrent.GUID:
				raise Exception("ERROR: torrent GUID does not match")
				self.kill()
				return

			#self.sock.setblocking(0)
			self.sock.settimeout(0.5)
			self.torrent.addPeer( self )

		except (socket.error, socket.timeout, Exception) as exception:
			#threadPrint( self, "HANDSHAKE FAIL", exception )
			self.kill()
			return






		


	def run (self):
		self.initFromTracker()
		
		while self.alive:
			# the receive half
			if self.recvSize == 0:
				self.checkForMessage()
	
			if self.recvSize > 0:
				ration = self.recvThrottle.getBandwidth( self.recvSize )
				
				try:
					data = self.sock.recv( ration )

					if len(data) > 0:
						self.resetTimeout()
						self.recvData += data
						self.recvSize -= len(data)
						self.recvTotal += len(data)
						#threadPrint( "\t", self, "sent", len(data), "bytes -> total", self.recvTotal )

						if self.recvSize == 0:
							self.handleMessage()
							self.checkForMessage()
					
				except socket.timeout:
					pass
				except socket.error as err:
					threadPrint( "!!!!!!!!!!!!!!!!!", self, err )
					break

			# the send half
			if len(self.sendData) == 0:
				if not self.sendQueue.empty():
					self.sendData = self.sendQueue.get()

			if len(self.sendData) > 0:
				ration = self.sendThrottle.getBandwidth( len(self.sendData) )

				try:
					sent = self.sock.send( self.sendData[:ration] )
					self.sendData = self.sendData[sent:]
					self.sendTotal += sent
					self.resetKeepAliveTimer()
					
				except socket.timeout:
					pass
				except socket.error as err:
					threadPrint( "!!!!!!!!!!!!!!!!!", self, err )
					break

			# if there's nothing going on, check back in a little while
			if self.recvSize == 0 and len(self.sendData) == 0 and self.sendQueue.empty():
				time.sleep(0.5)

		self.kill()








	def kill (self, sentByTimeout=False):
		if self.alive:
			threadPrint( self, "harakiri" )
			self.alive = False
			self.torrent.killPeer(self)
			self.torrent = None
			
			if self.sock != None:
				self.sock.close()
			self.sock = None

			if not sentByTimeout and self.timeout != None:
				self.timeout.cancel()
			self.timeout = None
			
			if self.keepAliveTimer != None:
				self.keepAliveTimer.cancel()
			self.keepAliveTimer = None


	def resetKeepAliveTimer (self):
		if self.keepAliveTimer != None:
			self.keepAliveTimer.cancel()
			
		if self.alive:
			self.keepAliveTimer = Timer( 90, self.sendKeepAlive )
			self.keepAliveTimer.start()


	def resetTimeout (self):
		if self.timeout != None:
			self.timeout.cancel()
				
		if self.alive:
			self.timeout = Timer( 150, self.kill, [True] )
			self.timeout.start()






# these are the send functions
	def setChoke (self, state):
		if self.choked != state:
			self.choked = state
			if self.choked:
				self.sendQueue.put( struct.pack("!LB", 1, CHOKE) )
				#threadPrint( "\t\tSEND", self, "choke" )
			else:
				self.sendQueue.put( struct.pack("!LB", 1, UNCHOKE) )
				#threadPrint( "\t\tSEND", self, "unchoke" ) 
		
	def updateInteresting (self):
		interest = False
		for piece in range(len(self.pieceComplete)):
			if self.pieceComplete[piece]:
				if not self.torrent.pieceComplete[piece]:
					interest = True
					break
	
		if self.interesting != interest:
			self.interesting = interest
			if self.interesting:
				self.sendQueue.put( struct.pack("!LB", 1, INTERESTED))
				#threadPrint( "\t\tSEND", self, "interested" )
			else:
				self.sendQueue.put( struct.pack("!LB", 1, NOT_INTERESTED))
				#threadPrint( "\t\tSEND", self, "not interested" )

	def sendKeepAlive (self):
		self.sendQueue.put( struct.pack("!L", 0))
		#threadPrint( "\t\tSEND", self, "keepAlive" )

	def sendHave (self, piece):
		if not self.pieceComplete[piece]:
			self.sendQueue.put( struct.pack("!LBL", 5, HAVE, piece))
			#threadPrint( "\t\tSEND", self, "have", piece )

	def sendBitfield (self):
		bitfield = self.torrent.bitfield
		layout = "!LB" + str(len(bitfield)) + "s"
		packedData = struct.pack(layout, 1 + len(bitfield), BITFIELD, bitfield)
		self.sendQueue.put( packedData )
		#threadPrint( "\t\tSEND", self, "bitfield" )

	def updateActiveRequests (self):
		if not self.choking_me:
			while len(self.activeRequests) < ACTIVE_REQUESTS_ALLOWED:
			
				if len(self.queuedRequests) > 0:
					request = self.queuedRequests.pop()
				else:
					request = self.torrent.nextFragmentRequest(self)
					if request == None:
						break
							
				self.sendRequest( request )
				
	def sendRequest (self, request):
		if request == None or request in self.queuedRequests or request in self.activeRequests:
			return

		if not self.choking_me and len(self.activeRequests) < ACTIVE_REQUESTS_ALLOWED:
			self.activeRequests.append( request )
			piece, offset, size = request
			packedData = struct.pack("!LBLLL", 13, REQUEST, piece, offset, size)
			self.sendQueue.put( packedData )
			#threadPrint( "\t\tSEND", self, "request", piece, offset, size )		

		else:
			self.queuedRequests.append( request )
				

	def cancelRequest (self, request):
		if request in self.activeRequests:
			self.activeRequests.remove( request )
			self.torrent.requestFailed( request )
			piece, offset, size = request
			packedData = struct.pack("!LBLLL", 13, CANCEL, piece, offset, size)
			self.sendQueue.put( packedData )
			#threadPrint( "\t\tSEND", self, "cancel", piece, offset, size )	

		elif request in self.queuedRequests:
			self.queuedRequests.remove( request )
			

	def sendFragment (self, piece, offset, fragment):
		layout = "!LBLL" + str(len(fragment)) + "s"
		packedData = struct.pack(layout, 9 + len(fragment), FRAGMENT, piece, offset, fragment)
		self.sendQueue.put( packedData )
		self.torrent.reportUpload( len(fragment) )
		self.fragmentsSent += 1
		#threadPrint( "\t\tSEND", self, "fragment", piece, offset )


# and these are the receive functions

	def checkForMessage (self):
		timeout = self.sock.gettimeout()
		self.sock.setblocking(1)
	
		try:
			data = ""
			while len(data) < 4:
				data = self.sock.recv(4 - len(data))
		except socket.error:
			self.kill()
			return

		self.recvSize = struct.unpack("!L", data)[0]
		if self.recvSize == 0:
			self.resetTimeout()

		self.sock.settimeout(timeout)
		#threadPrint( "\t\t\t\t", self, "sending a message of size", self.recvSize )

			
	def handleMessage (self):
		message = self.recvData
		self.recvData = ""
		
		messageType = struct.unpack("!B", message[0])[0]

		if messageType == CHOKE:
			#threadPrint( self, "choke" )
			self.choking_me = True
			
		elif messageType == UNCHOKE:
			#threadPrint( self, "unchoke" )
			self.choking_me = False
			self.updateActiveRequests()
			
		elif messageType == INTERESTED:
			#threadPrint( self, "interested" )
			self.interested_in_me = True
			
		elif messageType == NOT_INTERESTED:
			#threadPrint( self, "not interested" )
			self.interested_in_me = False

		elif messageType == BITFIELD:
			#threadPrint( self, "bitfield" )
			self.readBitfield( message[1:] )
			self.updateInteresting()
			self.updateActiveRequests()

			
		elif messageType == HAVE:
			index = struct.unpack("!L", message[1:5])[0]
			#threadPrint( self, "have", index )
			if not self.pieceComplete[index]:
				self.pieceComplete[index] = True
				self.torrent.incrementPiecePopulation(index)
			self.updateInteresting()
			self.updateActiveRequests()
			
		elif messageType == REQUEST:
			if not self.choked:
				piece, offset, length = struct.unpack("!LLL", message[1:13])
				#threadPrint( self, "request", piece, offset, length )
				fragment = self.torrent.getFragment( piece, offset, length )

				if fragment != None:
					self.sendFragment( piece, offset, fragment )

		# this feature is not yet implemented, and I'm not in a big hurry
		# in fact, I *CANT* implement it without a real outgoing request queue
		# right now, I just fire fragments out immediately. Not so cancellable.
		elif messageType == CANCEL:
			index, offset, length = struct.unpack("!LLL", message[1:13])
			#threadPrint( self, "cancel", index, offset, length )


		elif messageType == FRAGMENT:
			piece, offset = struct.unpack("!LL", message[1:9])
			fragment = message[9:]
			#threadPrint( self, "fragment", piece, offset, len(fragment) )
			self.torrent.receiveFragment( piece, offset, fragment )
			self.fragmentsReceived += 1

			for request in self.activeRequests:
				if piece == request[0] and offset == request[1]:
					self.activeRequests.remove( request )
			self.updateActiveRequests()


	def readBitfield (self, bitfield):
		# read the bitfield into a list
		bitfieldPieceComplete = []
		for byte in bitfield:
			for bit in range(8):
				if (ord(byte) >> bit) & 1:
					bitfieldPieceComplete.append( True )
				else:
					bitfieldPieceComplete.append( False )

		# copy the list across, updating the torrent as needed
		for i in range(self.torrent.pieceCount):
			if self.pieceComplete[i] and not bitfieldPieceComplete[i]:
				self.torrent.decrementPiecePopulation(i)
			elif not self.pieceComplete[i] and bitfieldPieceComplete[i]:
				self.torrent.incrementPiecePopulation(i)
			self.pieceComplete[i] = bitfieldPieceComplete[i]


	def __repr__ (self):
		return "<" + self.ip.rjust(15) + ", " + str(self.port).rjust(5) + ">"


			
